import cv2
import rclpy
from cv_bridge import CvBridge
from rclpy.node import Node
from sensor_msgs.msg import Image

# 将ros2 bag视频转化为mp4格式，注意修改要导出的话题
class ImageSubscriber(Node):
    def __init__(self, node_name):
        super().__init__(node_name)
        self.count = 0
        self.subscriber_ = self.create_subscription(Image, '/camera1/camera1/infra2/image_rect_raw', self.listener_callback, 10)
        self.cv_bridge = CvBridge()
        self.video_writer = None

        # 初始化视频写入器将在接收到第一帧时进行
        self.first_frame_received = False
        # register shutdown callback
        # self.create_shutdown_callback(self.on_shutdow)

    def listener_callback(self, data):
        try:
            image = self.cv_bridge.imgmsg_to_cv2(data, 'bgr8')
            if not self.first_frame_received:
                # 设置视频写入器，使用接收到的第一帧的尺寸
                height, width = image.shape[:2]
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                fps = 30  # 假设与发布者发布的速率一致，或根据需要调整
                self.video_writer = cv2.VideoWriter('output_video.mp4', fourcc, fps, (width, height))
                self.first_frame_received = True

            if rclpy.ok():
                #cv2.imshow("Received Image", image)
                self.video_writer.write(image)
                self.count = self.count + 1
                self.get_logger().info('image [%d/2193]' % self.count)
        except Exception as e:
            self.get_logger().error(f"Error processing image: {e}")

    def on_shutdown(self):
        self.get_logger().info("Shutting down subscriber...")
        if self.video_writer:
            self.video_writer.release()
            self.get_logger().info("shutting down video")
        cv2.destroyAllWindows()

# 确保在节点关闭时调用on_shutdown
def main(args=None):
    rclpy.init(args=args)
    node = ImageSubscriber('image_subscriber')
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.on_shutdown()  # 手动调用以确保资源释放
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
